import dotenv
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

dalle = OpenAIDALLEImageGenerationTool(
    name = "openai-dalle",
    api_wrapper = DallEAPIWrapper(model="dall-e-3")
)

llm = ChatOpenAI(model="gpt-4o-mini")
llm_with_tools = llm.bind_tools(tools=[dalle],tool_choice="openai-dalle")

chain = llm_with_tools|(lambda msg:msg.tool_calls[0]["args"])|dalle

# print(chain.invoke("帮我生成一张炒股票看大盘点照片"))

print(chain.invoke("帮我生成一张9:16尺寸，高端医院单人病房的照片"))